package xyz.thoughtset.viewer.models.openai;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.ai.zhipuai.ZhiPuAiChatOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.stereotype.Component;
import xyz.thoughtset.viewer.common.ai.model.entity.AiNode;
import xyz.thoughtset.viewer.common.ai.model.entity.ModelParam;
import xyz.thoughtset.viewer.common.ai.model.entity.purpose.ChatModelSetting;
import xyz.thoughtset.viewer.common.ai.model.factory.DefaultModelBuilder;

import java.util.Optional;

@Component
public class ZhipuChatBuilder extends DefaultModelBuilder {
    public ZhipuChatBuilder() {
        super("ZhipuAI");
    }


    @Override
    public ChatModel buildMode(AiNode node, ModelParam modelParam) {
        ZhiPuAiChatOptions.Builder builder = ZhiPuAiChatOptions.builder()
                .model(modelParam.getModel());
        ChatModelSetting settingObj = (ChatModelSetting) modelParam.getModelArgs();
        Optional.ofNullable(settingObj).ifPresent(setting -> {
            setOptionalValue(setting.getMaxTokens(), value -> builder.maxTokens(value.intValue()));
            setOptionalValue(setting.getTemperature(), builder::temperature);
        });
        ZhiPuAiChatOptions options = builder.build();
        ZhiPuAiApi api = new ZhiPuAiApi(node.getBaseUrl(), node.getApiKey());
        ChatModel chatModel = new ZhiPuAiChatModel(api, options);
        return chatModel;
    }
}
